import logging
import numpy as np
import faiss
import pickle
from pathlib import Path
from dataclasses import dataclass
from typing import Dict, Tuple, List, Any
from abc import ABC, abstractmethod
from sklearn.base import BaseEstimator
from sklearn.model_selection import train_test_split
import pandas as pd
from lazypredict.Supervised import LazyClassifier
import warnings
import sys

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)
warnings.filterwarnings('ignore')

@dataclass
class ModelResult:
    """模型结果数据类"""
    model_name: str
    accuracy: float
    balanced_accuracy: float
    roc_auc: float
    f1_score: float
    time_taken: float
    model: BaseEstimator

class DataProcessor:
    """数据处理类"""
    def __init__(self, train_folder: str, dest_folder: str):
        self.train_folder = Path(train_folder)
        self.dest_folder = Path(dest_folder)

    def load_data(self) -> Tuple[np.ndarray, np.ndarray]:
        """加载数据并进行预处理"""
        from util import createXY
        logger.info("开始数据加载...")
        X, y = createXY(str(self.train_folder), str(self.dest_folder))
        return self._preprocess_data(X, y)

    def _preprocess_data(self, X: np.ndarray, y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """数据预处理"""
        X = np.array(X).astype('float32')
        faiss.normalize_L2(X)
        y = np.array(y).astype('int32')
        return X, y

class ModelEvaluator(ABC):
    """模型评估抽象基类"""
    @abstractmethod
    def evaluate(self, X_train: np.ndarray, X_test: np.ndarray, 
                y_train: np.ndarray, y_test: np.ndarray) -> List[ModelResult]:
        pass

class LazyModelEvaluator(ModelEvaluator):
    """LazyPredict模型评估器"""
    def evaluate(self, X_train: np.ndarray, X_test: np.ndarray,
                y_train: np.ndarray, y_test: np.ndarray) -> List[ModelResult]:
        clf = LazyClassifier(verbose=0, ignore_warnings=True, custom_metric=None, random_state=42)
        results, _ = clf.fit(X_train, X_test, y_train, y_test)
        
        model_results = []
        for model_name, row in results.iterrows():
            result = ModelResult(
                model_name=model_name,
                accuracy=row['Accuracy'],
                balanced_accuracy=row['Balanced Accuracy'],
                roc_auc=row['ROC AUC'],
                f1_score=row['F1 Score'],
                time_taken=row['Time Taken'],
                model=clf.models.get(model_name)
            )
            model_results.append(result)
        
        return model_results

class ExperimentTracker:
    """实验跟踪器"""
    def __init__(self):
        self.results: List[ModelResult] = []
    
    def add_result(self, result: ModelResult):
        self.results.append(result)
    
    def get_best_model(self) -> ModelResult:
        return max(self.results, key=lambda x: x.f1_score)
    
    def format_results(self) -> pd.DataFrame:
        data = []
        for result in self.results:
            data.append({
                'Model': result.model_name,
                'Accuracy': result.accuracy,
                'Balanced Accuracy': result.balanced_accuracy,
                'ROC AUC': result.roc_auc,
                'F1 Score': result.f1_score,
                'Time Taken': result.time_taken
            })
        df = pd.DataFrame(data)
        return df.sort_values('Accuracy', ascending=False)

class ModelTrainingPipeline:
    """模型训练管道"""
    def __init__(self, data_processor: DataProcessor, 
                 evaluator: ModelEvaluator,
                 experiment_tracker: ExperimentTracker):
        self.data_processor = data_processor
        self.evaluator = evaluator
        self.experiment_tracker = experiment_tracker

    def run(self) -> Tuple[BaseEstimator, pd.DataFrame]:
        """执行训练管道"""
        try:
            # 加载和处理数据
            X, y = self.data_processor.load_data()
            logger.info("数据预处理完成")

            # 分割数据集
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=0.25, random_state=2023
            )
            logger.info(f"数据集分割完成: 训练集 {X_train.shape}, 测试集 {X_test.shape}")

            # 评估模型
            logger.info("开始模型评估...")
            model_results = self.evaluator.evaluate(X_train, X_test, y_train, y_test)
            
            # 跟踪结果
            for result in model_results:
                self.experiment_tracker.add_result(result)
            
            # 获取最佳模型
            best_result = self.experiment_tracker.get_best_model()
            logger.info(f"最佳模型: {best_result.model_name} (F1: {best_result.f1_score:.4f})")
            
            # 保存最佳模型
            self._save_model(best_result.model)
            
            # 格式化结果
            results_df = self.experiment_tracker.format_results()
            self._display_results(results_df)
            
            return best_result.model, results_df

        except Exception as e:
            logger.error(f"训练管道执行失败: {str(e)}")
            raise

    def _save_model(self, model: BaseEstimator):
        """保存模型"""
        logger.info("保存最佳模型...")
        with open('best_model.pkl', 'wb') as f:
            pickle.dump(model, f)
        logger.info("模型保存完成")

    def _display_results(self, results_df: pd.DataFrame):
        """显示结果"""
        print("\n模型评估结果:")
        print("-" * 80)
        pd.set_option('display.float_format', lambda x: '{:.4f}'.format(x))
        pd.set_option('display.max_columns', None)
        pd.set_option('display.width', None)
        print(results_df.to_string(index=False))

def main():
    """主函数"""
    try:
        # 初始化组件
        data_processor = DataProcessor(
            train_folder="../data/train",
            dest_folder="."
        )
        evaluator = LazyModelEvaluator()
        experiment_tracker = ExperimentTracker()
        
        # 创建并运行管道
        pipeline = ModelTrainingPipeline(
            data_processor=data_processor,
            evaluator=evaluator,
            experiment_tracker=experiment_tracker
        )
        
        best_model, results = pipeline.run()
        
    except KeyboardInterrupt:
        logger.info("\n程序被用户中断")
    except Exception as e:
        logger.error(f"程序执行失败: {str(e)}")
        raise

if __name__ == "__main__":
    main()